import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# -------------------------- 1. 生成飞机形状的3D体素 --------------------------
def create_plane_voxel(size=32):
    """生成飞机形状的3D体素数据（简化版飞机轮廓）"""
    voxel = torch.zeros(size, size, size)
    center = size // 2  # 体素中心坐标

    # 1. 机身（长方体）
    # 位置：x方向延伸，y/z居中
    body_x_start, body_x_end = center - 10, center + 5
    body_y_start, body_y_end = center - 2, center + 2
    body_z_start, body_z_end = center - 2, center + 2
    voxel[body_x_start:body_x_end, body_y_start:body_y_end, body_z_start:body_z_end] = 1.0

    # 2. 机翼（左右对称的长方体）
    # 位置：机身中部，y方向延伸
    wing_x_start, wing_x_end = center - 5, center + 1
    wing_y_start, wing_y_end = center - 8, center + 8
    wing_z_start, wing_z_end = center - 1, center + 1
    voxel[wing_x_start:wing_x_end, wing_y_start:wing_y_end, wing_z_start:wing_z_end] = 1.0

    # 3. 尾翼（垂直尾翼+水平尾翼）
    # 垂直尾翼（机身末端，z方向延伸）
    tail_vert_x_start, tail_vert_x_end = center + 3, center + 5
    tail_vert_y_start, tail_vert_y_end = center - 2, center + 2
    tail_vert_z_start, tail_vert_z_end = center - 4, center + 4
    voxel[tail_vert_x_start:tail_vert_x_end, tail_vert_y_start:tail_vert_y_end, tail_vert_z_start:tail_vert_z_end] = 1.0

    # 水平尾翼（机身末端，y方向延伸）
    tail_horiz_x_start, tail_horiz_x_end = center + 2, center + 4
    tail_horiz_y_start, tail_horiz_y_end = center - 4, center + 4
    tail_horiz_z_start, tail_horiz_z_end = center - 1, center + 1
    voxel[tail_horiz_x_start:tail_horiz_x_end, tail_horiz_y_start:tail_horiz_y_end, tail_horiz_z_start:tail_horiz_z_end] = 1.0

    return voxel


# -------------------------- 2. 3D模型可视化 --------------------------
def plot_voxel(voxel, title="3D Voxel Model"):
    """可视化3D体素模型"""
    fig = plt.figure(figsize=(8, 8))
    ax = fig.add_subplot(111, projection='3d')
    # 提取体素中“存在物体”的坐标
    x, y, z = np.where(voxel.numpy() > 0.5)
    ax.scatter(x, y, z, c='blue', marker='s', s=10)
    ax.set_title(title)
    plt.show()


# -------------------------- 3. 用PyTorch 3D网络学习生成飞机 --------------------------
class Plane3DGenerator(nn.Module):
    """3D生成网络：输入随机噪声，输出飞机形状的32×32×32体素"""
    def __init__(self, latent_dim=100):
        super(Plane3DGenerator, self).__init__()
        self.latent_dim = latent_dim
        # 3D转置卷积（上采样）：从噪声生成3D体素
        self.model = nn.Sequential(
            nn.ConvTranspose3d(latent_dim, 512, kernel_size=4, stride=1, padding=0),
            nn.ReLU(),
            nn.ConvTranspose3d(512, 256, kernel_size=4, stride=2, padding=1),
            nn.ReLU(),
            nn.ConvTranspose3d(256, 128, kernel_size=4, stride=2, padding=1),
            nn.ReLU(),
            nn.ConvTranspose3d(128, 1, kernel_size=4, stride=2, padding=1),
            nn.Sigmoid()  # 归一化到[0,1]
        )

    def forward(self, z):
        z = z.view(z.size(0), self.latent_dim, 1, 1, 1)
        voxel = self.model(z)
        return voxel.squeeze(1)


# -------------------------- 训练网络（学习生成飞机形状） --------------------------
def train_generator(generator, target_voxel, epochs=1000, lr=0.001):
    """训练生成网络，让其输出接近目标飞机体素"""
    criterion = nn.BCELoss()  # 二分类交叉熵损失（适配Sigmoid输出）
    optimizer = torch.optim.Adam(generator.parameters(), lr=lr)
    target_voxel = target_voxel.unsqueeze(0)  # 增加batch维度（batch_size=1）

    for epoch in range(epochs):
        optimizer.zero_grad()
        # 输入随机噪声
        noise = torch.randn(1, 100)
        generated_voxel = generator(noise)
        # 计算损失（生成结果与目标飞机的差异）
        loss = criterion(generated_voxel, target_voxel)
        loss.backward()
        optimizer.step()

        # 每100轮打印损失
        if (epoch + 1) % 100 == 0:
            print(f"Epoch {epoch+1}/{epochs}, Loss: {loss.item():.6f}")

    return generator


# -------------------------- 运行代码 --------------------------
if __name__ == "__main__":
    # 1. 生成飞机体素并可视化
    plane_voxel = create_plane_voxel()
    plot_voxel(plane_voxel, title="Plane Voxel Model")

    # 2. 训练生成网络，学习生成飞机形状
    generator = Plane3DGenerator(latent_dim=100)
    trained_generator = train_generator(generator, plane_voxel, epochs=1000)

    # 3. 用训练后的网络生成飞机模型并可视化
    noise = torch.randn(1, 100)
    generated_plane = trained_generator(noise)
    generated_plane = (generated_plane > 0.5).float()  # 二值化
    plot_voxel(generated_plane.squeeze(0), title="Generated Plane Voxel Model")